#!/usr/bin/env python 
# -*- coding: utf-8 -*-
# @Time    : 2023/4/19 13:10
# @Author  : WJQ
# @Site    : 
# @File    : mongo_handle.py
# @Software: PyCharm
import os, sys
import pymongo
from bson.objectid import ObjectId
from pymongo.collection import Collection


class MongoHandler(object):
    def __init__(self, db_name):
        mongo_client = pymongo.MongoClient(host='127.0.0.1', port=27017)
        # mongo 流量 db 名
        self.db_data = mongo_client[db_name]

    def insert_data(self, collection_name, item):
        db_collections = Collection(self.db_data, collection_name)
        return db_collections.insert_one(item)

    def is_data_exists(self, collection_name, data):
        collection = self.db_data[collection_name]
        # 使用数据中的 request 字段作为查询条件
        query = {"request": data["request"]}
        count = collection.count_documents(query)

        # 如果查询结果计数大于0，则表示存在相同的数据
        return count > 0

    def get_data_by_id(self, collection_name, data_id):
        try:
            collection = self.db_data[collection_name]
            data = collection.find_one({"_id": ObjectId(data_id)})
            return data
        except Exception as e:
            print(f"Error fetching data by ID: {e}")
            return None

    def get_all_data(self, collection_name):
        db_collections = Collection(self.db_data, collection_name)
        return list(db_collections.find())

# mongo = Handle_mongo_traffic_db()
